#pragma once

#include <ast/ast.h>
#include <dot/dot.h>
#include <common.h>
#include <unordered_map>

namespace ast {
  class NodeCompUnit;
  class NodeGlobalDecl;
  class NodeGlobalDefs;
  class NodeGlobalDef;
  class NodeConstDecl;
  class NodeConstDefs;
  class NodeConstDef;
  class NodeVarDecl;
  class NodeVarDefs;
  class NodeVarDef;
  class NodeArrayConst;
  class NodeArray;
  class NodeArrayInit;
  class NodeFuncDef;
  class NodeFuncDefArgs;
  class NodeFuncCall;
  class NodeBlock;
  class NodeBlockItems;
  class NodeStmt;
  class NodeStmtAssign;
  class NodeStmtReturn;
  class NodeStmtIf;
  class NodeStmtWhile;
  class NodeStmtBreak;
  class NodeStmtContinue;
  class NodeStmtFuncCall;
  class NodeExp;
  class NodeFuncArgs;
  class NodeCond;
}

class AstDot {
public:
  AstDot(std::ostream & out);
  ~AstDot();

  void visit(ast::NodeCompUnit & comp_unit);
  void visit(ast::NodeGlobalDecl & global_decl);
  void visit(ast::NodeGlobalDefs & global_defs);
  void visit(ast::NodeGlobalDef & global_def);
  void visit(ast::NodeConstDecl & const_decl);
  void visit(ast::NodeConstDefs & const_defs);
  void visit(ast::NodeConstDef & const_def);
  void visit(ast::NodeVarDecl & var_decl);
  void visit(ast::NodeVarDefs & var_defs);
  void visit(ast::NodeVarDef & var_def);
  void visit(ast::NodeArrayConst & array_const);
  void visit(ast::NodeArray & array);
  void visit(ast::NodeArrayInit & array_init);
  void visit(ast::NodeFuncDef & func_def);
  void visit(ast::NodeFuncDefArgs & func_def_args);
  void visit(ast::NodeFuncCall & func_call);
  void visit(ast::NodeBlock & block);
  void visit(ast::NodeBlockItems & block_items);
  void visit(ast::NodeStmt & stmt);
  void visit(ast::NodeStmtAssign & stmt_assign);
  void visit(ast::NodeStmtReturn & stmt_return);
  void visit(ast::NodeStmtIf & stmt_if);
  void visit(ast::NodeStmtWhile & stmt_while);
  void visit(ast::NodeStmtBreak & stmt_break);
  void visit(ast::NodeStmtContinue & stmt_continue);
  void visit(ast::NodeStmtFuncCall & stmt_func_call);
  void visit(ast::NodeExp & exp);
  void visit(ast::NodeFuncArgs & func_args);
  void visit(ast::NodeCond & cond);

  void write();

private:
  int id;
  Dot * dot;
  std::map<int, std::string> declare_list;
  std::unordered_multimap<int, std::tuple<int, std::string>> direction_list;

  void add_declare(int id, std::string & label);
  void add_direction(int id_src, int id_dst);
  void add_direction(int id_src, int id_dst, std::string label);
};
